from typing import List, Tuple
import time
import numpy as np
from utils.similarity import cosine_similar

# 近似最近邻查找
def search(features:np.ndarray, indices, invert_indices=None, topK=5):

    start_time = time.time()

    if invert_indices:
        # match centroid
        centroid_count = len(invert_indices)
        centroid_dists = []
        for i in range(centroid_count):
            dist = cosine_similar(invert_indices[i]['centroid'], features)
            centroid_dists.append((i, dist))
        centroid_dists.sort(key=lambda x: x[1], reverse=True)

        # 根据最接近的两个质心生成候选列表
        closest_label = centroid_dists[0][0]
        second_cloest_label = centroid_dists[1][0]
        candidate_list = invert_indices[closest_label]['ids'] + invert_indices[second_cloest_label]['ids']
    else:
        candidate_list = list(range(len(indices)))

    candidate_dists = []
    for i in range(len(candidate_list)):
        id = candidate_list[i]
        dist = cosine_similar(indices[id][0], features)
        
        candidate_dists.append((id, dist))
    candidate_dists.sort(key=lambda x: x[1], reverse=True)
    candidate_dists = candidate_dists[:topK]
    
    end_time = time.time()
    print(f"index search cost time: {end_time - start_time}s")

    ret = []
    for id, dist in candidate_dists:
        path = indices[id][1]
        ret.append((path, dist))
    return ret

def mixture_of_experts_search(features:List[np.ndarray], experts:List[Tuple], weights:List[float], topK=5):
    n_experts = len(experts)
    start_time = time.time()

    candidate_list = set()
    for i in range(n_experts):
        invert_indices = experts[i][1]
        if invert_indices:
            # match centroid
            centroid_count = len(invert_indices)
            centroid_dists = []
            for label in range(centroid_count):
                dist = cosine_similar(invert_indices[label]['centroid'], features[i])
                centroid_dists.append((label, dist))
            centroid_dists.sort(key=lambda x: x[1], reverse=True)

            # 根据最接近的两个质心生成候选列表
            closest_label = centroid_dists[0][0]
            second_cloest_label = centroid_dists[1][0]
            ids = invert_indices[closest_label]['ids'] + invert_indices[second_cloest_label]['ids']
            for id in ids:
                candidate_list.add(id)

    candidate_list = list(candidate_list)
    if len(candidate_list) == 0:
        candidate_list = list(range(len(experts[0][0])))

    candidate_dists = []
    for id in candidate_list:
        dist = 0
        # 混合不同特征的距离
        for i in range(n_experts):
            dist += cosine_similar(experts[i][0][id][0], features[i]) * weights[i]
        candidate_dists.append((id, dist))

    candidate_dists.sort(key=lambda x: x[1], reverse=True)
    candidate_dists = candidate_dists[:topK]
    
    end_time = time.time()
    print(f"index search cost time: {end_time - start_time}s")

    ret = []
    for id, dist in candidate_dists:
        path = experts[0][0][id][1]
        ret.append((path, dist))
    return ret